{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import seaborn as sns\n",
    "\n",
    "import sys\n",
    "sys.path.append(os.path.abspath('../'))\n",
    "from common import *\n",
    "import configs\n",
    "\n",
    "import json\n",
    "from benchmark import sum_stats\n",
    "import matplotlib.patches as mpatches\n",
    "import pandas as pd\n",
    "\n",
    "fig_name = 'benchmark'\n",
    "output_dir = join(FIGURE_DIR, fig_name)\n",
    "\n",
    "sns.set()\n",
    "sns.set_style('darkgrid', {'axes.edgecolor': 'black'})\n",
    "matplotlib.rcParams['axes.linewidth'] = 0.5\n",
    "sdfs = ['bunny_128', 'logo_256', 'shadowing_128']\n",
    "\n",
    "used_configs = [configs.OnlyShadingGrad(), configs.Warp(), configs.ConvolutionWarp2(), configs.ConvolutionWarp4(),\n",
    "                configs.ConvolutionWarp8(), configs.ConvolutionWarp16(), configs.ConvolutionWarp32()]\n",
    "all_stats = {}\n",
    "for sdf in sdfs:\n",
    "    all_stats[sdf] = {}\n",
    "    for config in used_configs:\n",
    "        stats_fn = join(output_dir, f'{sdf.split(\"_\")[0]}_{config.name}.json')\n",
    "        with open(stats_fn, 'r') as f:\n",
    "            all_stats[sdf][config.name] = json.load(f)\n",
    "\n",
    "# Average over all sdfs\n",
    "average_stats = {}\n",
    "for config in used_configs:\n",
    "    average_stats[config.name] = {}\n",
    "    for sdf in sdfs:\n",
    "        sum_stats(average_stats[config.name], all_stats[sdf][config.name], len(sdfs))\n",
    "\n",
    "renaming = [\n",
    "    'Only shading\\n gradients',\n",
    "    r'\\textbf{Ours}', 'Conv. \\n 2 samples',\n",
    "    'Conv. \\n 4 samples', 'Conv. \\n 8 samples',\n",
    "    'Conv. \\n 16 samples', 'Conv. \\n 32 samples'\n",
    "]\n",
    "\n",
    "for cfg, new_name in zip(used_configs, renaming):\n",
    "    average_stats[new_name] = average_stats.pop(cfg.name)\n",
    "\n",
    "df = pd.DataFrame.from_dict(average_stats).T\n",
    "df = df / 1000\n",
    "total_times = df.sum(axis=1).reset_index()\n",
    "total_times.columns = ['Technique', 'Total']\n",
    "\n",
    "y_offset = -0.1\n",
    "fontsize = 12\n",
    "base_size = 3\n",
    "n_rows = 2\n",
    "n_cols = 1\n",
    "fig = plt.figure(1, figsize=(COLUMN_WIDTH, 1.2 * COLUMN_WIDTH), constrained_layout=False)\n",
    "gs = fig.add_gridspec(n_rows, n_cols, wspace=0.025, hspace=0.29)\n",
    "\n",
    "ax = fig.add_subplot(gs[0])\n",
    "bar1 = sns.barplot(x=\"Technique\", y=\"Total\", data=total_times, color=sns.color_palette()[0], ax=ax)\n",
    "df = df.reset_index()\n",
    "df.columns = ['Technique', 'Primal rendering', 'Adjoint rendering']\n",
    "bar2 = sns.barplot(x=\"Technique\", y=\"Primal rendering\", data=df, color=sns.color_palette()[1], ax=ax)\n",
    "\n",
    "# Set up legend\n",
    "top_bar = mpatches.Patch(color=sns.color_palette()[0], label='Adjoint rendering')\n",
    "bottom_bar = mpatches.Patch(color=sns.color_palette()[1], label='Primal rendering')\n",
    "plt.legend(handles=[top_bar, bottom_bar])\n",
    "ax.set_ylabel('Time (s)', fontsize=fontsize)\n",
    "ax.set_xlabel('')\n",
    "ax.set_title(\"(a) Rendering time, $256^2$ pixels at 256 primal SPP and 64 adjoint SPP\", fontsize=fontsize)\n",
    "ax.bar_label(ax.containers[0], fmt='%.3f')\n",
    "ax.set_ylim(0, 4.75)\n",
    "\n",
    "# Second plot showing FSM timings\n",
    "stats_fn = join(output_dir, f'fsm.json')\n",
    "with open(stats_fn, 'r') as f:\n",
    "    fsm_stats = json.load(f)\n",
    "\n",
    "df = pd.DataFrame(fsm_stats)\n",
    "avg_times = df.mean(axis=1).reset_index()\n",
    "avg_times.columns = ['Resolution', 'Time']\n",
    "ax = fig.add_subplot(gs[1])\n",
    "bar1 = sns.barplot(x=\"Resolution\", y=\"Time\", data=avg_times, color=sns.color_palette()[0], ax=ax)\n",
    "v_offset = 0.02\n",
    "w_scale = 0.5\n",
    "ax.set_yscale('log')\n",
    "\n",
    "ax.bar_label(ax.containers[0], fmt='%.3f')\n",
    "ax.set_ylabel('Time (s)', fontsize=fontsize)\n",
    "ax.set_xlabel('')\n",
    "ax.set_title(\"(b) SDF redistancing time for different grid resolutions\",  fontsize=fontsize)\n",
    "ax.set_ylim(0, 2.7)\n",
    "ax.set_xlim(-0.5, 5.48)\n",
    "resolutions = avg_times['Resolution']\n",
    "bar1.set(xticklabels=[f'${r}^3$' for r in resolutions])\n",
    "print()\n",
    "\n",
    "# show the graph\n",
    "plt.margins(0, 0)\n",
    "# save_fig(fig_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "c8cc192903caa17681aa39d71202092ac11a526d37a1c4ad2948f13605924304"
  },
  "kernelspec": {
   "display_name": "Python 3.10.4 ('mi')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "c8cc192903caa17681aa39d71202092ac11a526d37a1c4ad2948f13605924304"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
